package mr

import (
	"fmt"
	"log"
	"net"
	"net/http"
	"net/rpc"
	"os"
	"sync"
	"time"
)

const (
	kMapStage    = 1
	kReduceStage = 2

	kMapTaskType    = 1
	kReduceTaskType = 2

	kNotDispatched = -1
	kFinished      = 0
)

type Coordinator struct {
	// Your definitions here.
	inputFiles              []string   // 输入文件，下标就是map task id
	mapTasksDispathcInfo    []int64    // mapTask 分配出去的时间
	intermediateFlies       [][]string // 中间文件，intermdeiate[i]表示reduce task i 的输入文件
	reduceTasksDispathcInfo []int64    // reduceTask 分配出去的时间
	stage                   int        // 所处阶段 1: map, 2 : reduce
	nReduce                 int        // reduce task的数量
	taskDoneNum             int        // map/reduce 已完成任务数量
	mutex                   sync.Mutex // 锁
}

// Your code here -- RPC handlers for the worker to call.
func (c *Coordinator) DispatchTask(arg DispatchTaskArgs, reply *DispatchTaskReply) error {
	now := time.Now().Unix()
	if c.stage == kMapStage {
		var taskFile string

		c.mutex.Lock()                      ////
		fmt.Println(c.mapTasksDispathcInfo) // for debug
		for i, time := range c.mapTasksDispathcInfo {
			if time == kNotDispatched {
				taskFile = c.inputFiles[i]
				c.mapTasksDispathcInfo[i] = now
				fmt.Println(now)
				break
			}
		}

		var taskId int = -1
		if taskFile != "" { // 存在可分配任务
			for i := range c.inputFiles {
				if c.inputFiles[i] == taskFile {
					taskId = i
					break
				}
			}
		}
		c.mutex.Unlock() ////

		reply.TaskType = kMapTaskType
		reply.TaskId = taskId
		reply.File = taskFile
		reply.NReduce = c.nReduce

	} else {
		// 请求一个reduce任务
		var reduceTaskId int = -1

		c.mutex.Lock() ////
		for i := range c.reduceTasksDispathcInfo {
			if c.reduceTasksDispathcInfo[i] == kNotDispatched {
				c.reduceTasksDispathcInfo[i] = now
				reduceTaskId = i
				break
			}
		}

		reply.TaskType = kReduceTaskType
		reply.TaskId = reduceTaskId
		if reduceTaskId != -1 {
			reply.File = c.intermediateFlies[reduceTaskId]
		} else {
			reply.File = nil
		}
		reply.NReduce = c.nReduce

		c.mutex.Unlock() ////
	}

	log.Println("dispatchTask, reply = ", reply)

	return nil
}

func (c *Coordinator) FinishTask(args FinishTaskArgs, reply *FinishTaskReply) error {
	log.Println("FinishTask start, args: ", args)
	taskId := args.TaskId
	if args.TaskType == kMapTaskType { // map任务完成
		c.mutex.Lock()                                   ////
		if c.mapTasksDispathcInfo[taskId] != kFinished { // 加个if从而只保留第一次task结果
			c.mapTasksDispathcInfo[taskId] = kFinished
			for reduceId := range args.IntermediateFiles {
				if args.IntermediateFiles[reduceId] != "" {
					c.intermediateFlies[reduceId] = append(c.intermediateFlies[reduceId], args.IntermediateFiles[reduceId])
				}
			}

			c.taskDoneNum++
			if c.taskDoneNum == len(c.inputFiles) {
				c.taskDoneNum = 0
				c.stage = kReduceStage
			}
		}
		c.mutex.Unlock() ////

	} else { // reduce任务完成
		c.mutex.Lock() ////
		if c.reduceTasksDispathcInfo[taskId] != kFinished {
			c.reduceTasksDispathcInfo[taskId] = kFinished
			c.taskDoneNum++
		}
		c.mutex.Unlock() ////
	}

	log.Println("FinishTash success, args:", args)

	return nil
}

//
// an example RPC handler.
//
// the RPC argument and reply types are defined in rpc.go.
//
func (c *Coordinator) Example(args *ExampleArgs, reply *ExampleReply) error {
	reply.Y = args.X + 1
	return nil
}

//
// start a thread that listens for RPCs from worker.go
//
func (c *Coordinator) server() {
	rpc.Register(c)
	rpc.HandleHTTP()
	//l, e := net.Listen("tcp", ":1234")
	sockname := coordinatorSock()
	os.Remove(sockname)
	l, e := net.Listen("unix", sockname)
	if e != nil {
		log.Fatal("listen error:", e)
	}
	go http.Serve(l, nil)
	log.Println("rpc server start")
}

//
// main/mrcoordinator.go calls Done() periodically to find out
// if the entire job has finished.
//
func (c *Coordinator) Done() bool {
	ret := false

	// Your code here.
	c.mutex.Lock() ////
	ret = c.stage == kReduceStage && c.taskDoneNum == c.nReduce
	log.Println("Done(): state:", c.stage, "taskDoneNum:", c.taskDoneNum)
	c.mutex.Unlock() ////

	return ret
}

// 检查故障任务（另开一个线程执行）
func (c *Coordinator) checkCrash() {
	for !c.Done() {
		log.Println("checkCrash")

		c.mutex.Lock() ////
		now := time.Now().Unix()
		if c.stage == kMapStage {
			for i, dispatchTime := range c.mapTasksDispathcInfo {
				if dispatchTime > 0 && now-dispatchTime > 10 {
					c.mapTasksDispathcInfo[i] = kNotDispatched
				}
			}
		} else {
			for i, dispatchTime := range c.reduceTasksDispathcInfo {
				if dispatchTime > 0 && now-dispatchTime > 10 {
					c.reduceTasksDispathcInfo[i] = kNotDispatched
				}
			}
		}
		c.mutex.Unlock() ////

		time.Sleep(time.Second)
	}
}

//
// create a Coordinator.
// main/mrcoordinator.go calls this function.
// nReduce is the number of reduce tasks to use.
//
func MakeCoordinator(files []string, nReduce int) *Coordinator {
	c := Coordinator{}

	// Your code here.

	// init
	c.inputFiles = files
	c.mapTasksDispathcInfo = make([]int64, len(files))
	for i := range c.mapTasksDispathcInfo {
		c.mapTasksDispathcInfo[i] = kNotDispatched
	}
	c.reduceTasksDispathcInfo = make([]int64, nReduce)
	for i := range c.reduceTasksDispathcInfo {
		c.reduceTasksDispathcInfo[i] = kNotDispatched
	}
	c.intermediateFlies = make([][]string, nReduce)
	c.stage = kMapStage
	c.nReduce = nReduce
	c.taskDoneNum = 0

	log.Println("map task num:", len(c.inputFiles)) // for debug

	// 开个线程检查执行太慢的task
	go c.checkCrash()

	// 创个文件夹，存放中间文件
	os.RemoveAll("intermediate")
	os.MkdirAll("intermediate", 0777)

	c.server()
	return &c
}
